import xarray as xr
import numpy as np
import os
import glob
from tqdm import tqdm
import matplotlib.pyplot as plt


def create_mask(data):
    """创建有效数据区域的掩码"""
    # 找出至少有一个时间步有数据的网格点
    mask = ~np.isnan(data).all(axis=0)
    return mask


def crop_to_valid_region(data, mask):
    """裁剪到有效数据区域"""
    # 找出有效区域的边界
    rows = np.any(mask, axis=1)
    cols = np.any(mask, axis=0)

    # 获取边界索引
    row_start, row_end = np.where(rows)[0][[0, -1]]
    col_start, col_end = np.where(cols)[0][[0, -1]]

    # 裁剪数据
    cropped_data = data[:, row_start:row_end + 1, col_start:col_end + 1]
    cropped_mask = mask[row_start:row_end + 1, col_start:col_end + 1]

    return cropped_data, cropped_mask, (row_start, row_end, col_start, col_end)


def analyze_data_distribution(data, var_name):
    """分析数据分布"""
    # 创建有效数据掩码
    mask = create_mask(data)

    # 计算有效数据比例
    valid_ratio = np.mean(mask)
    print(f"\nValid data analysis for {var_name}:")
    print(f"Valid data ratio: {valid_ratio:.2%}")

    # 绘制数据分布图
    plt.figure(figsize=(15, 5))

    # 有效数据区域
    plt.subplot(131)
    plt.imshow(mask, cmap='binary')
    plt.title(f'{var_name} Valid Data Region')
    plt.colorbar()

    # 数据平均值分布
    plt.subplot(132)
    mean_data = np.nanmean(data, axis=0)
    plt.imshow(mean_data, cmap='viridis')
    plt.title(f'{var_name} Mean Value Distribution')
    plt.colorbar()

    # 数据标准差分布
    plt.subplot(133)
    std_data = np.nanstd(data, axis=0)
    plt.imshow(std_data, cmap='viridis')
    plt.title(f'{var_name} Standard Deviation Distribution')
    plt.colorbar()

    plt.tight_layout()
    plt.savefig(f'{var_name}_distribution.png')
    plt.close()


def process_nc_file(nc_file, output_dir='.'):
    """处理单个nc文件"""
    print(f"Processing {nc_file}...")

    # 打开nc文件
    ds = xr.open_dataset(nc_file)

    # 获取变量名
    var_name = list(ds.data_vars.keys())[0]
    print(f"Variable name: {var_name}")

    # 获取数据
    data = ds[var_name].values

    # 打印原始数据信息
    print("\nOriginal data information:")
    print(f"Data shape: {data.shape}")
    print(f"Data type: {data.dtype}")
    print(f"Min value: {np.nanmin(data)}")
    print(f"Max value: {np.nanmax(data)}")
    print(f"Mean value: {np.nanmean(data)}")
    print(f"Missing values: {np.isnan(data).sum()}")

    # 分析数据分布
    analyze_data_distribution(data, var_name)

    # 创建掩码并裁剪数据
    mask = create_mask(data)
    cropped_data, cropped_mask, bounds = crop_to_valid_region(data, mask)

    # 打印裁剪后数据信息
    print("\nCropped data information:")
    print(f"New shape: {cropped_data.shape}")
    print(f"Bounds: {bounds}")
    print(f"Min value: {np.nanmin(cropped_data)}")
    print(f"Max value: {np.nanmax(cropped_data)}")
    print(f"Mean value: {np.nanmean(cropped_data)}")
    print(f"Missing values: {np.isnan(cropped_data).sum()}")

    # 保存原始数据
    original_file = os.path.join(output_dir, f"{var_name}_original.npy")
    np.save(original_file, data)
    print(f"\nSaved original data to {original_file}")

    # 保存裁剪后的数据
    cropped_file = os.path.join(output_dir, f"{var_name}_cropped.npy")
    np.save(cropped_file, cropped_data)
    print(f"Saved cropped data to {cropped_file}")

    # 保存掩码
    mask_file = os.path.join(output_dir, f"{var_name}_mask.npy")
    np.save(mask_file, cropped_mask)
    print(f"Saved mask to {mask_file}")

    # 保存边界信息
    bounds_file = os.path.join(output_dir, f"{var_name}_bounds.npy")
    np.save(bounds_file, bounds)
    print(f"Saved bounds to {bounds_file}")

    return data, cropped_data, cropped_mask, bounds


def main():
    # 设置输入输出目录
    input_dir = '.'
    output_dir = 'processed_nc'
    os.makedirs(output_dir, exist_ok=True)

    # 获取所有nc文件
    nc_files = glob.glob(os.path.join(input_dir, '*.nc'))
    print(f"Found {len(nc_files)} nc files")

    # 处理每个文件
    all_original_data = {}
    all_cropped_data = {}
    all_masks = {}
    all_bounds = {}

    for nc_file in tqdm(nc_files):
        # 从文件名获取变量名
        var_name = os.path.basename(nc_file).split('_')[0]

        # 处理文件
        original_data, cropped_data, mask, bounds = process_nc_file(nc_file, output_dir)

        # 存储数据
        all_original_data[var_name] = original_data
        all_cropped_data[var_name] = cropped_data
        all_masks[var_name] = mask
        all_bounds[var_name] = bounds

    # 保存合并后的数据
    for var_name in all_original_data.keys():
        # 保存原始数据
        original_file = os.path.join(output_dir, f"{var_name}_combined_original.npy")
        np.save(original_file, all_original_data[var_name])
        print(f"\nCombined original data for {var_name}:")
        print(f"Final shape: {all_original_data[var_name].shape}")
        print(f"Missing values: {np.isnan(all_original_data[var_name]).sum()}")
        print(f"Saved to {original_file}")

        # 保存裁剪数据
        cropped_file = os.path.join(output_dir, f"{var_name}_combined_cropped.npy")
        np.save(cropped_file, all_cropped_data[var_name])
        print(f"\nCombined cropped data for {var_name}:")
        print(f"Final shape: {all_cropped_data[var_name].shape}")
        print(f"Missing values: {np.isnan(all_cropped_data[var_name]).sum()}")
        print(f"Saved to {cropped_file}")

        # 保存掩码
        mask_file = os.path.join(output_dir, f"{var_name}_combined_mask.npy")
        np.save(mask_file, all_masks[var_name])
        print(f"Saved combined mask to {mask_file}")

        # 保存边界
        bounds_file = os.path.join(output_dir, f"{var_name}_combined_bounds.npy")
        np.save(bounds_file, all_bounds[var_name])
        print(f"Saved combined bounds to {bounds_file}")


if __name__ == '__main__':
    main() 